"""

x1,y1为恶意代码的标签和特征

x2,y2为正常代码的标签和特征

要做的：
1循环读多个文件
2放样本

完成



接下来要做的：
先尝试深度学习



"""

import os
import re

import matplotlib.pyplot as plt
from sklearn.feature_extraction.text import CountVectorizer
from sklearn import model_selection
from sklearn.datasets import load_iris
from sklearn import tree
import pydotplus
import numpy as np
from sklearn.neighbors import KNeighborsClassifier



def load_one_flle(filename):
    x=[]
    with open(filename) as f:
        line=f.readline()
        line=line.strip('\n')
    return line




def load_files(rootdir,num):
    x=[]
    y=[]
    list = os.listdir(rootdir)
    for i in range(0, len(list)):
        path = os.path.join(rootdir, list[i])
        if os.path.isfile(path):
            x.append(load_one_flle(path))
            y.append(num)
    return x,y




def dirlist(path, allfile):
    filelist = os.listdir(path)

    for filename in filelist:
        filepath = os.path.join(path, filename)
        if os.path.isdir(filepath):
            dirlist(filepath, allfile)
        else:
            allfile.append(filepath)
    return allfile




if __name__ == '__main__':
    x1, y1 = load_files("normal/",0)
    x2, y2 = load_files('virus/',1)
    print(type(x1))
    print(y1)
    print(x2)
    print(y2)
    x = x1+x2
    y = y1+y2
    print(x)
    print(y)
    vectorizer = CountVectorizer(min_df=1)
    x=vectorizer.fit_transform(x)
    x=x.toarray()
    print (x)
    print (y)
    clf = KNeighborsClassifier(n_neighbors=3)
    clf.fit(x,y)
    scores=model_selection.cross_val_score(clf, x, y, n_jobs=-1, cv=4)
    y_pred = clf.predict([[55,65,51,33,67,35,46,45,4,59,22,44,76,43,74,20,0,42,15,70,27,40,29,3,11,78,26,28,53,37,17,39,71,16,14,62,66,8,63,38,12,56,57,72,73,75,77,5,30,36,58,21,69,64,9,24,52,2,6,7,25,1,68,60,18,54,23,19,13]])
    print(y_pred)
    y_pred2 = clf.predict([[67,51,65,33,35,46,45,4,76,22,44,59,43,74,47,70,20,42,15,0,27,40,29,11,78,26,28,37,17,39,71,16,14,62,66,38,12,75,73,72,57,77,30,36,69,58,21,64,9,52,2,7,25,1,68,18,23,19,13,32,31,10,41,34,0,0,0,0,0]])
    print(y_pred2)
    print(scores)
    print(np.mean(scores))
